import cv2
# 显示结果
import fitz
import numpy as np
from PIL import Image

from paddleocr import PaddleOCR, draw_ocr

# test_img_path = './imgs/1.png'
font_path = '../doc/fonts/simfang.ttf'


# test_pdf_path = './xxx.pdf'


def ocr_img_test(test_img_path):
    # Paddleocr目前支持的多语言语种可以通过修改lang参数进行切换
    # 例如`ch`, `en`, `fr`, `german`, `korean`, `japan`
    # use_angle_cls=True, , cls=True
    # ocr = PaddleOCR(lang="ch", use_angle_cls=True, det_model_dir='D:/ocr/inf/ch_PP-OCRv3_det_infer',
    #                 rec_model_dir='D:/ocr/output/rec_ppocr_v3_inf')
    ocr = PaddleOCR(lang="ch", rec_model_dir='D:/ocr/output/v3/best_accuracy',
                    rec_char_dict_path='D:/workspace/python/train_data/train_data_configs/words/result_all.txt',
                    use_gpu=False)
    # ocr = PaddleOCR(lang="ch", pretrained_model='D:/ocr/output/best_barkup1/best_accuracy') #无效
    # ocr = PaddleOCR(lang="ch")
    result = ocr.ocr(test_img_path)
    # result = ocr.ocr(test_img_path)
    for idx in range(len(result)):
        res = result[idx]
        for line in res:
            print(line)
    return result


pil('D:/ocr/ocr_resources/train_data/char_enum_by_font_no_dot/images_1/000010430.jpg')


# ocr_img_test('D:/ocr/ocr_resouces/train_data/train_data_customer/crop_img/1_crop_0.jpg')
# ocr_img_test('D:/ocr/ocr_resouces/train_data/train_data_customer/1.png')
# ocr_img_test('D:/ocr/ocr_resouces/train_data_version2/train_data_spz/00046879.jpg')  # 生僻字有问题
# ocr_img_test('D:/ocr/ocr_resouces/train_data_version2/train_data_jj/train_images/img_370.jpg')#街景有问题
# ocr_img_test('D:/ocr/ocr_resouces/train_data_version2/train_data_zy/00062522.jpg')
# ocr_img_test('D:/ocr/ocr_resouces/train_data_version2/train_data_jj/train_images/img_127.jpg')#街景有问题


def pil(test_img_path):
    result = ocr_img_test(test_img_path)
    # 显示结果
    # 如果本地没有simfang.ttf，可以在doc/fonts目录下下载
    result = result[0]
    image = Image.open(test_img_path).convert('RGB')
    boxes = [line[0] for line in result]
    txts = [line[1][0] for line in result]
    scores = [line[1][1] for line in result]
    im_show = draw_ocr(image, boxes, txts, scores, font_path)
    im_show = Image.fromarray(im_show)
    im_show.save('result.jpg')


def ocr_pdf_test(test_pdf_path):
    # Paddleocr目前支持的多语言语种可以通过修改lang参数进行切换
    # 例如`ch`, `en`, `fr`, `german`, `korean`, `japan`
    ocr = PaddleOCR(use_angle_cls=True, lang="ch",
                    page_num=2)  # need to run only once to download and load model into memory
    result = ocr.ocr(test_pdf_path, cls=True)
    for idx in range(len(result)):
        res = result[idx]
        for line in res:
            print(line)
    return result


def ocr_pdf_img(test_pdf_path):
    imgs = []
    result = ocr_pdf_test()
    with fitz.open(test_pdf_path) as pdf:
        for pg in range(0, pdf.pageCount):
            page = pdf[pg]
            mat = fitz.Matrix(2, 2)
            pm = page.getPixmap(matrix=mat, alpha=False)
            # if width or height > 2000 pixels, don't enlarge the image
            if pm.width > 2000 or pm.height > 2000:
                pm = page.getPixmap(matrix=fitz.Matrix(1, 1), alpha=False)

            img = Image.frombytes("RGB", [pm.width, pm.height], pm.samples)
            img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
            imgs.append(img)
    for idx in range(len(result)):
        res = result[idx]
        image = imgs[idx]
        boxes = [line[0] for line in res]
        txts = [line[1][0] for line in res]
        scores = [line[1][1] for line in res]
        im_show = draw_ocr(image, boxes, txts, scores, font_path)
        im_show = Image.fromarray(im_show)
        im_show.save('result_page_{}.jpg'.format(idx))

# pil()
# ocr_pdf_img()
